import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.cm import ScalarMappable
import matplotlib.colors as mcolors

# 读取点云数据
point_cloud = o3d.io.read_point_cloud("temp/point_cloud_after_icp.ply")
bim_mesh = o3d.io.read_triangle_mesh("BIM/OBJ/qiaoduntai.obj")

# 将三角网格转换为点云（用于拟合平面）
bim_point_cloud = bim_mesh.sample_points_uniformly(number_of_points=50000000)
bim_point_cloud2 = bim_mesh.sample_points_uniformly(number_of_points=100000)

distances = np.array(point_cloud.compute_point_cloud_distance(bim_point_cloud))

# 设置阈值
threshold = 0.01

# 拆分点云为两部分：低于阈值的点和超过阈值的点
below_point = o3d.geometry.PointCloud()
above_point = o3d.geometry.PointCloud()

# 存储低于阈值和超过阈值的点
below_points = []
above_points = []
above_distances = []

for i, dist in enumerate(distances):
    if dist <= threshold:
        below_points.append(np.asarray(point_cloud.points)[i])
    else:
        above_points.append(np.asarray(point_cloud.points)[i])
        above_distances.append(dist)

# 将拆分后的点云数据赋值给对应的点云对象
if below_points:
    below_point.points = o3d.utility.Vector3dVector(np.array(below_points))
    below_point.paint_uniform_color([0.5, 0.5, 0.5])  # 灰色

if above_points:
    above_point.points = o3d.utility.Vector3dVector(np.array(above_points))

    # 对超过阈值的点进行颜色映射
    above_distances = np.array(above_distances)
    sorted_indices = np.argsort(above_distances)[::-1]  # 从大到小排序
    sorted_distances = above_distances[sorted_indices]

    min_error = sorted_distances.min()
    max_error = sorted_distances.max()

    # 获取超过阈值的点的颜色
    above_colors = []
    for dist in sorted_distances:
        normalized_error = (dist - min_error) / (max_error - min_error)  # 归一化
        color = plt.cm.coolwarm(normalized_error)[:3]  # 从红色到蓝色
        above_colors.append(color)

    # 设置超过阈值点的颜色
    above_point.colors = o3d.utility.Vector3dVector(np.array(above_colors)[sorted_indices])

# 为BIM点云设置灰色
bim_point_cloud2.paint_uniform_color([0.5, 0.5, 0.5])

# 计算并打印误差信息
if above_points:
    average_error = np.mean(sorted_distances)
    print(f"Average Registration Error: {average_error:.4f} m")
    print(f"Minimum Error: {min_error:.4f} m")
    print(f"Maximum Error: {max_error:.4f} m")
    print(f"Number of Points Exceeding Threshold: {len(above_points)}")

    # 绘制误差距离和颜色关系的可视化
    fig, ax = plt.subplots(figsize=(4, 8))
    norm = mcolors.Normalize(vmin=min_error, vmax=max_error)
    cmap = plt.cm.coolwarm
    cbar = plt.colorbar(
        ScalarMappable(norm=norm, cmap=cmap),
        ax=ax,
        label='Registration Error Distance (m)',
        orientation='vertical'
    )
    num_ticks = 5
    tick_locations = np.linspace(min_error, max_error, num_ticks)
    cbar.set_ticks(tick_locations)
    cbar.set_ticklabels([f'{x:.4f}' for x in tick_locations])
    ax.set_title('Error Distance to Color Mapping')
    cbar.ax.set_ylabel('Registration Error Distance (m)', rotation=270, labelpad=15)
    plt.tight_layout()
    plt.savefig("error_distribution.png")
    plt.show()
    plt.close(fig)

# 可视化
o3d.visualization.draw_geometries(
    [above_point,bim_point_cloud2],  # 同时显示误差点和BIM点云
    window_name="Point Cloud with Error Heatmap",
    point_show_normal=False
)

# 合并点云并保存
if below_point.has_points() and above_point.has_points():
    combined_pcd = below_point + above_point
elif below_point.has_points():
    combined_pcd = below_point
else:
    combined_pcd = above_point

print("误差结果点云保存")
o3d.io.write_point_cloud("Processed/output.ply", combined_pcd)
